import cv2
import numpy as np
import os
import re

class MangaPanelExtractor:
    def __init__(self):
        self.min_panel_size = 3000  # 最小分镜面积
        self.black_threshold = 80  # 黑色线条的阈值
        self.line_thickness_range = (2, 20)  # 分割线的粗细范围
        self.min_line_length_ratio = 0.5  # 最小线条长度占图像宽度/高度的比例
        self.debug = True

    def enhance_edges(self, image):
        """增强图像边缘"""
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用Canny边缘检测
        edges = cv2.Canny(gray, 50, 150)
        
        # 使用形态学操作加粗边缘
        kernel = np.ones((3,3), np.uint8)
        dilated = cv2.dilate(edges, kernel, iterations=1)
        
        return dilated

    def detect_lines(self, image):
        """检测黑色分割线，只保留长直线"""
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        height, width = gray.shape[:2]
        
        # 使用自适应阈值处理找出黑色区域
        binary = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY_INV, 25, 5
        )
        
        # 创建一个空白图像用于绘制检测到的线条
        line_image = np.zeros_like(binary)
        
        # 使用霍夫直线检测
        lines = cv2.HoughLinesP(
            binary,
            rho=1,
            theta=np.pi/180,
            threshold=100,
            minLineLength=min(width, height) * self.min_line_length_ratio,  # 最小线长
            maxLineGap=20  # 最大间隙
        )
        
        if lines is not None:
            for line in lines:
                x1, y1, x2, y2 = line[0]
                # 计算线条长度
                length = np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
                # 计算线条角度
                angle = np.abs(np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi)
                
                # 只保留接近水平或垂直的线条
                is_horizontal = (angle < 5) or (angle > 175)
                is_vertical = (85 < angle < 95)
                
                # 检查线条是否足够长
                is_long_enough = False
                if is_horizontal:
                    is_long_enough = length > width * self.min_line_length_ratio
                elif is_vertical:
                    is_long_enough = length > height * self.min_line_length_ratio
                
                # 如果是长直线，则绘制到line_image上
                if (is_horizontal or is_vertical) and is_long_enough:
                    cv2.line(line_image, (x1, y1), (x2, y2), 255, self.line_thickness_range[0])
        
        # 使用形态学操作连接断开的线条
        kernel_v = np.ones((self.line_thickness_range[1], 1), np.uint8)
        kernel_h = np.ones((1, self.line_thickness_range[1]), np.uint8)
        
        # 分别处理垂直和水平线
        vertical = cv2.morphologyEx(line_image, cv2.MORPH_CLOSE, kernel_v)
        horizontal = cv2.morphologyEx(line_image, cv2.MORPH_CLOSE, kernel_h)
        
        # 合并垂直和水平线
        combined = cv2.bitwise_or(vertical, horizontal)
        
        # 去除小的噪点
        kernel = np.ones((3,3), np.uint8)
        cleaned = cv2.morphologyEx(combined, cv2.MORPH_OPEN, kernel, iterations=1)
        
        return cleaned

    def find_panels(self, line_image, original_image):
        """查找分镜区域"""
        # 反转图像使分镜区域为白色
        inverted = cv2.bitwise_not(line_image)
        
        # 查找轮廓
        contours, _ = cv2.findContours(
            inverted, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        
        height, width = original_image.shape[:2]
        total_area = height * width
        min_area = total_area * 0.05  # 增加最小面积比例到5%
        max_area = total_area * 0.95  # 最大面积为总面积的95%
        
        valid_panels = []
        for contour in contours:
            area = cv2.contourArea(contour)
            if area < min_area or area > max_area:
                continue
            
            # 获取边界框
            x, y, w, h = cv2.boundingRect(contour)
            
            # 检查宽高比
            aspect_ratio = w / float(h)
            if aspect_ratio < 0.3 or aspect_ratio > 3:  # 调整宽高比范围
                continue
            
            # 检查区域内的内容
            roi = original_image[y:y+h, x:x+w]
            gray_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
            
            # 计算非白色像素的比例
            non_white_ratio = np.sum(gray_roi < 250) / float(w * h)
            if non_white_ratio < 0.1:  # 如果内容太少，跳过
                continue
            
            valid_panels.append((x, y, w, h))
        
        return valid_panels

    def sort_panels(self, panels):
        """按照从上到下，从左到右的顺序排序分镜"""
        # 计算每个分镜的中心y坐标
        panel_centers = [(p, p[1] + p[3]/2) for p in panels]
        
        # 按y坐标分组（允许一定的误差）
        y_tolerance = 50  # 同一行的y坐标允许的最大差异
        rows = []
        current_row = []
        last_y = None
        
        # 按y坐标排序
        sorted_centers = sorted(panel_centers, key=lambda x: x[1])
        
        for panel, center_y in sorted_centers:
            if last_y is None or abs(center_y - last_y) > y_tolerance:
                if current_row:
                    # 对当前行按x坐标排序
                    current_row.sort(key=lambda p: p[0])
                    rows.append(current_row)
                current_row = [panel]
                last_y = center_y
            else:
                current_row.append(panel)
        
        # 添加最后一行
        if current_row:
            current_row.sort(key=lambda p: p[0])
            rows.append(current_row)
        
        # 合并所有行
        sorted_panels = []
        for row in rows:
            sorted_panels.extend(row)
        
        return sorted_panels

    def extract_panels(self, image_path, output_dir):
        """提取漫画分镜"""
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        print(f"\n处理图片: {image_path}")
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图片: {image_path}")
        
        # 获取图片尺寸
        height, width = image.shape[:2]
        print(f"图片尺寸: {width}x{height}")
        
        # 检测分割线
        line_image = self.detect_lines(image)
        
        # 保存线条检测结果用于调试
        cv2.imwrite(os.path.join(output_dir, "lines.jpg"), line_image)
        
        # 查找分镜
        panels = self.find_panels(line_image, image)
        
        # 对分镜进行排序
        sorted_panels = self.sort_panels(panels)
        
        # 在原图上绘制检测结果
        debug_image = image.copy()
        
        # 提取并保存分镜
        for i, (x, y, w, h) in enumerate(sorted_panels):
            # 添加小边距
            margin = 5
            x = max(0, x - margin)
            y = max(0, y - margin)
            w = min(width - x, w + 2*margin)
            h = min(height - y, h + 2*margin)
            
            # 提取分镜
            panel = image[y:y+h, x:x+w]
            
            # 在调试图像上绘制边框和编号
            cv2.rectangle(debug_image, (x, y), (x+w, y+h), (0, 255, 0), 2)
            cv2.putText(debug_image, str(i+1), (x+10, y+30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            # 保存分镜
            output_path = os.path.join(output_dir, f"panel_{i+1:02d}.jpg")
            cv2.imwrite(output_path, panel)
            print(f"保存分镜 {i+1:02d}: {output_path}")
        
        # 保存调试图像
        cv2.imwrite(os.path.join(output_dir, "detected_panels.jpg"), debug_image)
        
        return len(sorted_panels)

def natural_sort_key(s):
    """用于实现文件夹的自然排序"""
    return [int(text) if text.isdigit() else text.lower()
            for text in re.split('([0-9]+)', str(s))]

def main():
    # 创建实例
    extractor = MangaPanelExtractor()
    
    # 确保输入和输出目录存在
    input_dir = "pic"
    output_base_dir = "extracted_panels"
    
    if not os.path.exists(input_dir):
        os.makedirs(input_dir)
        print(f"创建输入目录: {input_dir}")
    
    if not os.path.exists(output_base_dir):
        os.makedirs(output_base_dir)
        print(f"创建输出目录: {output_base_dir}")
    
    # 获取所有图片文件
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.jpg', '.png'))]
    
    print(f"\n找到 {len(image_files)} 个图片文件")
    
    if not image_files:
        print(f"\n错误：在 {input_dir} 目录中没有找到图片文件！")
        return
    
    # 进行自然排序
    image_files.sort(key=natural_sort_key)
    
    # 处理每个图片文件
    for image_file in image_files:
        print(f"\n处理文件: {image_file}")
        input_path = os.path.join(input_dir, image_file)
        
        # 创建输出目录
        output_dir = os.path.join(output_base_dir, os.path.splitext(image_file)[0])
        
        try:
            # 如果输出目录已存在，先删除
            if os.path.exists(output_dir):
                import shutil
                shutil.rmtree(output_dir)
                print(f"清理已存在的输出目录: {output_dir}")
            
            # 处理图片
            num_panels = extractor.extract_panels(input_path, output_dir)
            print(f"成功从 {image_file} 中提取了 {num_panels} 个分镜到 {output_dir}")
            
        except Exception as e:
            print(f"处理 {image_file} 时出错: {str(e)}")
            continue

if __name__ == "__main__":
    main() 